import multiprocessing
import multiprocessing.shared_memory
import time

import numpy as np
import scipy
import PIL.Image

def worker(i,shm1_name,shm2_name,ns,i_start,i_end):
    shm1 = multiprocessing.shared_memory.SharedMemory(name=shm1_name)
    shm2 = multiprocessing.shared_memory.SharedMemory(name=shm2_name)
    
    sxy = np.ndarray((ns,2), dtype=np.float64, buffer=shm1.buf)
    us = np.ndarray((ns,1), dtype=np.float64, buffer=shm2.buf)
    
    ncam = 2

    N_sphere = 2
    sphere_list = np.zeros((N_sphere,4)) #x,y,z,r
    sphere_list[0,:] = np.array([0,0,-10,2])
    sphere_list[1,:] = np.array([2,0,-5,1])

    N_light = 1
    light_list = np.zeros((N_light,3)) #x,y,z
    light_list[0,:] = np.array([0,-10,-10])

    def get_sphere_normal(point, collided_sphere_id):
        if point.ndim != 1:
            raise ValueError('123')
            
        normal = point - sphere_list[collided_sphere_id,0:3]
        return normal/np.linalg.norm(normal)

    def get_brightness(point,normal,view_direction=None):
        if point.ndim != 1 or normal.ndim != 1:
            raise ValueError('123')
            
        brightness = 0.2
        highlight = 0
        
        costheta = 0
        costheta_view = 0
        
        h = None
        
        for i in range(N_light):
            light_ray = light_list[i,:] - point
            
            costheta = light_ray[None,:]@normal[:,None]
            costheta = costheta/(np.linalg.norm(light_ray)*np.linalg.norm(normal))
            costheta = float(costheta[0,0])
            
            if costheta > 0:
                brightness += costheta
            
            if costheta > 0 and view_direction is not None:
                if view_direction.ndim != 1:
                    raise ValueError('123')
                
                h = (-view_direction/np.linalg.norm(view_direction)+light_ray/np.linalg.norm(light_ray))/2
                costheta_view = h[None,:]@normal[:,None]
                costheta_view = costheta_view/(np.linalg.norm(h)*np.linalg.norm(normal))
                costheta_view = float(costheta_view[0,0])
                
                if costheta_view > 0:
                    highlight = costheta_view**32
                    #print(highlight)
                    brightness += highlight
                    
            if brightness >= 1:
                break

        brightness = min(max(float(brightness), 0), 1)
        return brightness

    def is_ray_collide_with_this_sphere(point,direction, sphere_id):
        if point.ndim != 1 or direction.ndim != 1:
            raise ValueError('123')
            
        is_collided = False
        collided_coord = None
        
        A = np.sum(direction**2)
        B = (2*direction[None,:]@(point[:,None] - sphere_list[sphere_id:sphere_id+1,0:3].T))[0,0]
        C = np.sum((point[None,:] - sphere_list[sphere_id:sphere_id+1,0:3]) ** 2) - sphere_list[sphere_id,3]**2
        D = B**2 - 4*A*C
        
        is_collided = D > 0
        
        if is_collided:
            t = np.zeros((2))
            t[0] = (-B+np.sqrt(D))/(2*A)
            t[1] = (-B-np.sqrt(D))/(2*A)
            
            _temp_collided_coord = direction[None,:]*t[:,None]+point[None,:] #(2*3)
            
            if _temp_collided_coord[0,2] > _temp_collided_coord[1,2]:
                collided_coord = _temp_collided_coord[0,:]
            else:
                collided_coord = _temp_collided_coord[1,:]

        return is_collided, collided_coord

    def is_ray_collide_with_sphere(point,direction,use_Zbuffer=True):
        if point.ndim != 1 or direction.ndim != 1:
            raise ValueError('123')
        
        is_collided = False
        collided_sphere_id = 0
        collided_coord = np.zeros((3))
        
        z_buffer = -9999
        
        for i in range(N_sphere):
            _is_collided, _collided_coord = is_ray_collide_with_this_sphere(point,direction,i)
            
            if _is_collided and (_collided_coord[2] > z_buffer or use_Zbuffer == False):
                is_collided = True
                collided_sphere_id = i
                collided_coord = _collided_coord
                z_buffer = float(_collided_coord[2])
                
            if is_collided and use_Zbuffer == False:
                break
                
        return is_collided,collided_coord,collided_sphere_id
            
    for i in range(i_start,i_end): 
        xs = sxy[i,0]
        ys = sxy[i,1]

        is_collided, collided_coord, collided_sphere_id = is_ray_collide_with_sphere(np.array([0,0,ncam]),np.array([xs,ys,-ncam]))
        
        if is_collided:
            norm = get_sphere_normal(collided_coord, collided_sphere_id)
            brightness = get_brightness(collided_coord, norm, np.array([xs,ys,-ncam]))
            us[i,0] = brightness
    
    shm1.close()
    shm2.close()
    return

if __name__ == '__main__':
    
    dx = 0.0025
    L = 1

    xs,ys = np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx))
    sxy = np.column_stack((xs.flatten(),ys.flatten()))
    del xs,ys
    
    ns = sxy.shape[0]
    
    us = np.zeros((ns,1))
    
    shm1 = multiprocessing.shared_memory.SharedMemory(create=True, size=sxy.nbytes)
    shm2 = multiprocessing.shared_memory.SharedMemory(create=True, size=us.nbytes)
    shared_sxy = np.ndarray((ns,2), dtype=np.float64, buffer=shm1.buf)
    shared_us = np.ndarray((ns,1), dtype=np.float64, buffer=shm2.buf)
    shared_sxy[:,:] = sxy[:,:]
    shared_us[:,:] = 0
    
    print(ns)
    
    num_proc = 4
    process = []
    i_start = []
    i_end = []
    
    start_time = time.time()
    
    for i in range(num_proc):
        work_range = int(ns/num_proc)
        i_start.append(work_range*i)
        i_end.append(work_range*(i+1) if i < num_proc - 1 else ns)
        
        p = multiprocessing.Process(target=worker,args=(i,shm1.name,shm2.name,ns,i_start[i],i_end[i]))
        process.append(p)
        p.start()
        
    print(i_start)
    print(i_end)
    
    for p in process:
        p.join()
    
    us = shared_us.copy()
    image = np.array(us.reshape((int(np.sqrt(ns)),int(np.sqrt(ns))))*255, dtype = np.uint8)
    
    shm1.close()
    shm1.unlink()
    shm2.close()
    shm2.unlink()
    
    end_time = time.time()
    print(end_time-start_time)
    
    print(us.shape)
    print(us.max())
    image = PIL.Image.fromarray(image)
    image.save('output_image.png')
    image.show()
